import tensorflow as tf
import os
'''使用说明：
1.将train.tfrecord, validation.tfrecord分别放到两个不同的目录，目录下不要有其它文件
2.将把将路径写到变量train_data_path，validation_data_path中
3.将labels.txt的路径赋值给label_path
4.填好输出路径train_image_path和validation_image_path
5.运行脚本那可'''
dic = {}
num = 0
train_data_path = 'your train tfrecord path'
validation_data_path = 'your validation tfrecord path'
train_image_path = 'your output train image path'
validation_image_path = 'your output validation image path'
label_path = 'your label file path'

with open(label_path, 'rb') as f:
    for line in f:
        line = line.decode('utf-8').strip()
        label = line.split(':')[-1]
        dic[num] = label
        num+=1

train_files = tf.gfile.Glob(train_data_path)
filename_queue = tf.train.string_input_producer(train_files)
# 创建TFRecordReader实例
reader = tf.TFRecordReader()
# 从tfrecord中读取数据
_, serialize_example = reader.read(filename_queue)
features = tf.parse_single_example(serialize_example, features={
    'image/encoded': tf.FixedLenFeature([], tf.string),
    'image/width': tf.FixedLenFeature([], tf.int64),
    'image/height': tf.FixedLenFeature([], tf.int64),
    'image/format': tf.FixedLenFeature([], tf.string),
    'image/class/label': tf.FixedLenFeature([], tf.int64),
})
images = features['image/encoded']
widths = features['image/width']
heights = features['image/height']
formats = features['image/format']
labels = features['image/class/label']
with tf.Session() as sess:
    coord = tf.train.Coordinator()
    threads = tf.train.start_queue_runners(sess=sess, coord=coord)
    i = 0
    for record in tf.python_io.tf_record_iterator(filename_queue):
        i += 1
        image, label, width, height = sess.run([images, labels, widths, heights])
        with open(os.path.join(train_image_path,'{}_label_{}.jpg'.format(i, label)), 'wb') as f:
            f.write(image)
    print("train image num:%d" % i)
    coord.request_stop()
    coord.join(threads)

validation_files = tf.gfile.Glob(validation_data_path)
filename_queue = tf.train.string_input_producer(validation_files)
# 创建TFRecordReader实例
reader = tf.TFRecordReader()
# 从tfrecord中读取数据
_, serialize_example = reader.read(filename_queue)
features = tf.parse_single_example(serialize_example, features={
    'image/encoded': tf.FixedLenFeature([], tf.string),
    'image/width': tf.FixedLenFeature([], tf.int64),
    'image/height': tf.FixedLenFeature([], tf.int64),
    'image/format': tf.FixedLenFeature([], tf.string),
    'image/class/label': tf.FixedLenFeature([], tf.int64),
})
images = features['image/encoded']
widths = features['image/width']
heights = features['image/height']
formats = features['image/format']
labels = features['image/class/label']
with tf.Session() as sess:
    coord = tf.train.Coordinator()
    threads = tf.train.start_queue_runners(sess=sess, coord=coord)
    i = 0
    for record in tf.python_io.tf_record_iterator(filename_queue):
        i += 1
        image, label, width, height = sess.run([images, labels, widths, heights])
        with open(os.path.join(validation_image_path,'{}_label_{}.jpg'.format(i, label)), 'wb') as f:
            f.write(image)
    print("validation image num:%d" % i)
    coord.request_stop()
    coord.join(threads)
